import ray
import numpy as np

ray.init(ignore_reinit_error=True)


class SecretShareBase:

    def __init__(self,
                 row_count=10,
                 scale=1 << 10,
                 mod=1 << 64,
                 federation=None,
                 table_wrapper=None):
        self.row_count = row_count
        self.scale = scale
        self.mod = mod
        # self.keys = [1, 2, 4, 100]
        self.a_matrix = np.asarray([1] * self.row_count).reshape(
            self.row_count, 1)
        self.b_matrix = np.asarray([2] * self.row_count).reshape(
            self.row_count, 1)
        self.c_matrix = np.asarray([4] * self.row_count).reshape(
            self.row_count, 1)
        self.r_matrix = np.asarray([100] * self.row_count).reshape(
            self.row_count, 1)
        self.transfer_variable = SecretShareTransferVariable()
        if federation:
            RuntimeInstance.FEDERATION = federation
        if table_wrapper:
            RuntimeInstance.TABLE_WRAPPER = table_wrapper

    def split_series(self, series, r):
        return self.format_matrix(series) - self.format_matrix(
            r), self.format_matrix(r)

    def format_matrix(self, data):
        return np.mod(np.multiply(data, self.scale), self.mod)

    def add_value(self, a, b):

        def _fn(row):
            for i in range(len(row)):
                if row[i] > self.mod / 2:
                    row[i] -= self.mod / 2
            return row

        agg_sum = np.add(a, b)
        agg_sum = agg_sum % self.mod
        agg_sum = np.apply_along_axis(_fn, 0, agg_sum)
        agg_sum = agg_sum / (self.scale * self.scale)
        return agg_sum

    def rust_split_series(self, series, r):
        return self.rust_format_matrix(series) - self.rust_format_matrix(
            r), self.rust_format_matrix(r)

    def rust_format_matrix(self, data):
        return np.multiply(data, self.scale)

    def rust_add_value(self, a, b):
        agg_sum = np.add(a, b)
        agg_sum = agg_sum / (self.scale * self.scale)
        return agg_sum


class LeaderSecretShareMaster(SecretShareBase):

    def multiple_party_matrix_rustspeedup(self, leader_matrix, leader_f_count,
                                          partner_f_count, suffix,
                                          concurrent_slice_size, idx=-1):
        # shape = leader_matrix.shape
        # batch_list = np.array_split(leader_matrix,
        #                             shape[0] // concurrent_slice_size + 1)
        # batch_list_len = len(batch_list)
        # LOGGER.info(f"leader_matrix shape is {shape}")
        # LOGGER.info(
        #     f"batch_list_len: {batch_list_len}, slice_size:{concurrent_slice_size}"
        # )
        a, b, c, r = 1, 2, 4, 0 #2**32
        shape = partner_f_count
        a = a*np.ones(shape)
        b = b*np.ones(shape)
        c = c*np.ones(shape)
        r = r*np.ones(shape)
        # x0, _ = self.rust_split_series(leader_matrix, r)
        lm = leader_matrix * self.scale
        for i in range(shape[1]-1):
            lm = np.c_[lm, leader_matrix * self.scale]
        x0 = lm - r
        # LOGGER.info(f"leader_matrix is {leader_matrix}")
        # y0 = self.rust_format_matrix(np.array(r))
        y0 = r
        u0 = x0 + a
        v0 = y0 + b
        self.transfer_variable.r_value.remote((u0, v0),
                                              suffix=('leader_u_v', ) + suffix,
                                              role=consts.PROVIDER,
                                              idx=idx)

        partner_u_v = \
            self.transfer_variable.secret_pair.get(idx=idx, suffix=('partner_u_v',)+suffix)
        u1, v1 = partner_u_v[0], partner_u_v[1]
        u = (u0 + u1)
        v = (v0 + v1)
        # LOGGER.info(f"u is {u}")
        # self.transfer_variable.r_value.remote((u, v), suffix=('final_u_v',)+suffix, role=consts.PROVIDER, idx=-1)
        fed_matrix = -b * u - a * v + c  # * self.scale
        # LOGGER.info(f"fed_matrix is {fed_matrix}")
        fed_sum = np.sum(fed_matrix, axis=0)
        return fed_sum

    def multiple_party_matrix_rustspeedup2(self, leader_matrix, leader_f_count,
                                          partner_f_count, suffix,
                                          concurrent_slice_size, idx_li):
        a, b, c, r = 1, 2, 4, 0
        shape = partner_f_count
        a = a*np.ones(shape)
        b = b*np.ones(shape)
        c = c*np.ones(shape)
        r = r*np.ones(shape)
        lm = leader_matrix * self.scale
        for i in range(shape[1]-1):
            lm = np.c_[lm, leader_matrix * self.scale]
        x0 = lm - r
        y0 = r
        u0 = x0 + a
        v0 = y0 + b
        self.transfer_variable.r_value.remote((u0, v0),
                                              suffix=('leader_u_v', ) + suffix,
                                              role=consts.PROVIDER,
                                              idx=-1)
        for idx in idx_li:
            partner_u_v = \
                self.transfer_variable.secret_pair.get(idx=idx, suffix=('partner_u_v',)+suffix)
            u1, v1 = partner_u_v[0], partner_u_v[1]
            u = (u0 + u1)
            v = (v0 + v1)
            fed_matrix = -b * u - a * v + c
            fed_sum = np.sum(fed_matrix, axis=0)

            self.transfer_variable.z0.remote(
                fed_sum,
                role=consts.PROVIDER,
                suffix=('z0',)+suffix,
                idx=idx)


    def multiple_party_matrix(self, leader_matrix):
        shape = leader_matrix.shape
        self.transfer_variable.promoter_col_count.remote(shape[1],
                                                         idx=0,
                                                         role=consts.PROVIDER)
        provider_col_count = self.transfer_variable.provider_col_count.get(
            idx=-1)[0]
        fed_value_list = []
        for i in range(shape[1]):
            rst, futures = [], []
            for j in range(provider_col_count):
                worker = LeaderSecretShareWorker.remote(
                    row_count=self.row_count,
                    federation=RuntimeInstance.FEDERATION,
                    table_wrapper=RuntimeInstance.TABLE_WRAPPER)
                future = worker.multiple_series.remote(leader_matrix[:, i],
                                                       suffix=f"{i}_{j}")
                futures.append(future)
            results = ray.get(futures)
            for fed_v_list in results:
                fed_value = fed_v_list.sum() % self.mod
                rst.append(fed_value)
            # for future in futures:
            #     fed_v_list = ray.get(future)
            #     fed_value = fed_v_list.sum() % self.mod
            #     rst.append(fed_value)
            fed_value_list.append(rst)

        return fed_value_list


@ray.remote
class LeaderSecretShareWorker(SecretShareBase):

    def multiple_series(self, series, suffix="suffix") -> np.ndarray:
        x0, x1 = self.split_series(series, self.r_matrix)
        y1 = self.format_matrix(np.array(self.r_matrix))
        v0 = self.transfer_variable.v0.get(idx=-1, suffix=suffix)[0]
        u1 = self.transfer_variable.u1.get(idx=-1, suffix=suffix)[0]
        u0 = (x0 + self.a_matrix) % self.mod
        v1 = (y1 + self.b_matrix) % self.mod
        u = (u0 + u1) % self.mod
        v = (v0 + v1) % self.mod
        self.transfer_variable.u.remote(u,
                                        idx=0,
                                        suffix=suffix,
                                        role=consts.PROVIDER)
        self.transfer_variable.v.remote(v,
                                        idx=0,
                                        suffix=suffix,
                                        role=consts.PROVIDER)
        z0 = self.c_matrix * self.scale - self.b_matrix * u - self.a_matrix * v
        return z0


class PartnerSecretShareMaster(SecretShareBase):

    def multiple_party_matrix_rustspeedup(self, partner_matrix, leader_f_count,
                                          partner_f_count, suffix,
                                          concurrent_slice_size):
        # shape = partner_matrix.shape
        # batch_list = np.array_split(partner_matrix,
        #                             shape[0] // concurrent_slice_size + 1)
        # batch_list_len = len(batch_list)
        # LOGGER.info(
        #     f"batch_list_len: {batch_list_len}, slice_size:{concurrent_slice_size}"
        # )
        a, b, c, r = 1, 2, 4, 0 #2**32
        shape = partner_matrix.shape
        a = a*np.ones(shape)
        b = b*np.ones(shape)
        c = c*np.ones(shape)
        r = r*np.ones(shape)
        # _, y1 = self.rust_split_series(partner_matrix, r)
        y1 = partner_matrix * self.scale - r
        LOGGER.info(f"partner_matrix is {partner_matrix.shape}")
        # LOGGER.info(f"y1 is {y1}")
        # x1 = self.rust_format_matrix(np.array(r))
        x1 = r
        # LOGGER.info(f"x1 is {x1}")
        v1 = (y1 + b)
        u1 = (x1 + a)
        self.transfer_variable.secret_pair.remote(
            (v1, u1),
            suffix=('partner_u_v', ) + suffix,
            role=consts.PROMOTER,
            idx=-1)
        leader_u_v = \
            self.transfer_variable.r_value.get(idx=0, suffix=('leader_u_v',)+suffix)
        u0, v0 = leader_u_v[0], leader_u_v[1]
        u = (u0 + u1)
        v = (v0 + v1)
        # LOGGER.info(f"u is {u}")
        # final_u_v = \
        # self.transfer_variable.r_value.get(idx=0, suffix=('final_u_v',)+suffix)
        # u, v = final_u_v[0], final_u_v[1]
        fed_matrix = u * v - b * u - a * v + c  # * self.scale
        # LOGGER.info(f"fed_matrix is {fed_matrix}")
        fed_sum = np.sum(fed_matrix, axis=0)
        return fed_sum

    def multiple_party_matrix_rustspeedup2(self, partner_matrix, leader_f_count,
                                          partner_f_count, suffix,
                                          concurrent_slice_size):
        a, b, c, r = 1, 2, 4, 0
        shape = partner_matrix.shape
        a = a*np.ones(shape)
        b = b*np.ones(shape)
        c = c*np.ones(shape)
        r = r*np.ones(shape)
        y1 = partner_matrix * self.scale - r
        x1 = r
        v1 = (y1 + b)
        u1 = (x1 + a)
        self.transfer_variable.secret_pair.remote(
            (v1, u1),
            suffix=('partner_u_v', ) + suffix,
            role=consts.PROMOTER,
            idx=-1)
        leader_u_v = \
            self.transfer_variable.r_value.get(idx=0, suffix=('leader_u_v',)+suffix)
        u0, v0 = leader_u_v[0], leader_u_v[1]
        u = (u0 + u1)
        v = (v0 + v1)
        fed_matrix = u * v - b * u - a * v + c
        fed_sum = np.sum(fed_matrix, axis=0)

        z0 = self.transfer_variable.z0.get(idx=0, suffix=("z0",)+suffix)

        z = self.rust_add_value(z0, fed_sum)
        return z

    def multiple_party_matrix(self, partner_matrix):
        shape = partner_matrix.shape
        self.transfer_variable.provider_col_count.remote(shape[1],
                                                         idx=0,
                                                         role=consts.PROMOTER)
        promoter_col_count = self.transfer_variable.promoter_col_count.get(
            idx=-1)[0]
        fed_value_list = []
        for i in range(promoter_col_count):
            rst, futures = [], []
            for j in range(shape[1]):
                worker = PartnerSecretShareWorker.remote(
                    row_count=self.row_count,
                    federation=RuntimeInstance.FEDERATION,
                    table_wrapper=RuntimeInstance.TABLE_WRAPPER)
                future = worker.multiple_series.remote(partner_matrix[:, j],
                                                       suffix=f"{i}_{j}")
                futures.append(future)

            results = ray.get(futures)
            for fed_v_list in results:
                fed_value = fed_v_list.sum() % self.mod
                rst.append(fed_value)
            fed_value_list.append(rst)
            # for future in futures:
            #     fed_v_list = ray.get(future)
            #     fed_value = fed_v_list.sum() % self.mod
            #     rst.append(fed_value)
            # fed_value_list.append(rst)
        return fed_value_list


@ray.remote
class PartnerSecretShareWorker(SecretShareBase):

    def multiple_series(self, series, suffix="suffix"):
        y0, y1 = self.split_series(series, self.r_matrix)
        x1 = self.format_matrix(np.array(self.r_matrix))
        v0 = (y0 + self.b_matrix) % self.mod
        u1 = (x1 + self.a_matrix) % self.mod
        self.transfer_variable.v0.remote(v0,
                                         idx=0,
                                         suffix=suffix,
                                         role=consts.PROMOTER)
        self.transfer_variable.u1.remote(u1,
                                         idx=0,
                                         suffix=suffix,
                                         role=consts.PROMOTER)
        u = self.transfer_variable.u.get(idx=-1, suffix=suffix)[0]
        v = self.transfer_variable.v.get(idx=-1, suffix=suffix)[0]
        z1 = u * v - self.b_matrix * u - self.a_matrix * v + self.c_matrix * self.scale
        return z1
